#include "Graphics/PiccoloVKPipeline.h"

#include <Graphics/PiccoloVKDevice.h>
#include <Graphics/PiccoloVKRenderPass.h>

#include "PiccoloFileUtil.h"

namespace Piccolo
{
    //-----------------------------------------------------PiccoloVKPipelineLayout---------------------------------------------------------

    PiccoloVKPipelineLayout::PiccoloVKPipelineLayout(PiccoloVKDevice *device, const std::string &vertexShaderFile, const std::string &fragShaderFile, const ShaderLayout &shaderLayout)
        : mDevice(device)
    {
        // 编译shader
        CALL_VK(CreateShaderModule(vertexShaderFile + ".spv", &mVertexShaderModule));
        CALL_VK(CreateShaderModule(fragShaderFile + ".spv", &mFragShaderModule));

        // 管线布局
        VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .setLayoutCount = static_cast<uint32_t>(shaderLayout.descriptorSetLayouts.size()),
            .pSetLayouts = shaderLayout.descriptorSetLayouts.data(),
            .pushConstantRangeCount = static_cast<uint32_t>(shaderLayout.pushConstantRanges.size()),
            .pPushConstantRanges = shaderLayout.pushConstantRanges.data()
        };

        CALL_VK(vkCreatePipelineLayout(mDevice->GetLogicDevice(), &pipelineLayoutCreateInfo, nullptr, &mPipelineLayout));
    }

    PiccoloVKPipelineLayout::~PiccoloVKPipelineLayout()
    {
        VK_D(ShaderModule, mDevice->GetLogicDevice(), mVertexShaderModule);
        VK_D(ShaderModule, mDevice->GetLogicDevice(), mFragShaderModule);
        VK_D(PipelineLayout, mDevice->GetLogicDevice(), mPipelineLayout);
    }

    VkResult PiccoloVKPipelineLayout::CreateShaderModule(const std::string &shaderFile, VkShaderModule*shaderModule)
    {
        std::vector<char> content = ReadCharArrayFromFile(shaderFile);

        VkShaderModuleCreateInfo shaderModuleCreateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .codeSize = static_cast<uint32_t>(content.size()),
            .pCode = reinterpret_cast<const uint32_t*>(content.data())
        };

        return vkCreateShaderModule(mDevice->GetLogicDevice(), &shaderModuleCreateInfo, nullptr, shaderModule);
    }

    //-----------------------------------------------------PiccoloVKPipeline---------------------------------------------------------

    PiccoloVKPipeline::PiccoloVKPipeline(PiccoloVKDevice *device, PiccoloVKRenderPass *renderPass, PiccoloVKPipelineLayout *pipelineLayout)
        : mDevice(device), mRenderPass(renderPass), mPipelineLayout(pipelineLayout)
    {
    }

    PiccoloVKPipeline::~PiccoloVKPipeline()
    {
        VK_D(Pipeline, mDevice->GetLogicDevice(), mPipeline);
    }

    void PiccoloVKPipeline::Create()
    {
        // shader
        VkPipelineShaderStageCreateInfo shaderStageInfo[] =
        {
            {
                .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
                .pNext = nullptr,
                .flags = 0,
                .stage = VK_SHADER_STAGE_VERTEX_BIT,
                .module = mPipelineLayout->GetVertexShaderModule(),
                .pName = "main",
                .pSpecializationInfo = nullptr
            },
            {
                .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
                .pNext = nullptr,
                .flags = 0,
                .stage = VK_SHADER_STAGE_FRAGMENT_BIT,
                .module = mPipelineLayout->GetFragShaderModule(),
                .pName = "main",
                .pSpecializationInfo = nullptr
            }
        };

        // 顶点输入
        VkPipelineVertexInputStateCreateInfo vertexInputStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .vertexBindingDescriptionCount = static_cast<uint32_t>(mPipelineConfig.vertexInputState.vertexBindingDescriptions.size()),
            .pVertexBindingDescriptions = mPipelineConfig.vertexInputState.vertexBindingDescriptions.data(),
            .vertexAttributeDescriptionCount = static_cast<uint32_t>(mPipelineConfig.vertexInputState.vertexAttributes.size()),
            .pVertexAttributeDescriptions = mPipelineConfig.vertexInputState.vertexAttributes.data()
        };

        VkPipelineInputAssemblyStateCreateInfo inputAssemblyStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .topology = mPipelineConfig.inputAssemblyState.topology,
            .primitiveRestartEnable = mPipelineConfig.inputAssemblyState.primitiveRestartEnable
        };

        // 视口
        VkViewport defaultViewport =
        {
            .x = 0,
            .y = 0,
            .width = 100,
            .height = 100,
            .minDepth = 0,
            .maxDepth = 1
        };
        VkRect2D defaultScissor =
        {
            .offset = {0, 0},
            .extent = {100, 100}
        };
        VkPipelineViewportStateCreateInfo viewportStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .viewportCount = 1,
            .pViewports = &defaultViewport,
            .scissorCount = 1,
            .pScissors = &defaultScissor
        };

        // 光栅化
        VkPipelineRasterizationStateCreateInfo rasterizationStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .depthClampEnable = mPipelineConfig.rasterizationState.depthClampEnable,
            .rasterizerDiscardEnable = mPipelineConfig.rasterizationState.rasterizerDiscardEnable,
            .polygonMode = mPipelineConfig.rasterizationState.polygonMode,
            .cullMode = mPipelineConfig.rasterizationState.cullMode,
            .frontFace = mPipelineConfig.rasterizationState.frontFace,
            .depthBiasEnable = mPipelineConfig.rasterizationState.depthBiasEnable,
            .depthBiasConstantFactor = mPipelineConfig.rasterizationState.depthBiasConstantFactor,
            .depthBiasClamp = mPipelineConfig.rasterizationState.depthBiasClamp,
            .depthBiasSlopeFactor = mPipelineConfig.rasterizationState.depthBiasSlopeFactor,
            .lineWidth = mPipelineConfig.rasterizationState.lineWidth
        };

        // 超采样
        VkPipelineMultisampleStateCreateInfo multisampleStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .rasterizationSamples = mPipelineConfig.multisampleState.rasterizationSamples,
            .sampleShadingEnable = mPipelineConfig.multisampleState.sampleShadingEnable,
            .minSampleShading = mPipelineConfig.multisampleState.minSampleShading,
            .pSampleMask = nullptr,
            .alphaToCoverageEnable = VK_FALSE,
            .alphaToOneEnable = VK_FALSE
        };

        // 深度模板
        VkPipelineDepthStencilStateCreateInfo depthStencilStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .depthTestEnable = mPipelineConfig.depthStencilState.depthTestEnable,
            .depthWriteEnable = mPipelineConfig.depthStencilState.depthWriteEnable,
            .depthCompareOp = mPipelineConfig.depthStencilState.depthCompareOp,
            .depthBoundsTestEnable = mPipelineConfig.depthStencilState.depthBoundsTestEnable,
            .stencilTestEnable = mPipelineConfig.depthStencilState.stencilTestEnable,
            .front = {},
            .back = {},
            .minDepthBounds = 0.f,
            .maxDepthBounds = 0.f
        };

        // 颜色混合
        VkPipelineColorBlendStateCreateInfo colorBlendStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .logicOpEnable = VK_FALSE,
            .logicOp = VK_LOGIC_OP_CLEAR,
            .attachmentCount = 1,
            .pAttachments = &mPipelineConfig.colorBlendAttachmentState
        };
        colorBlendStateInfo.blendConstants[0] = colorBlendStateInfo.blendConstants[1] = colorBlendStateInfo.blendConstants[2] = colorBlendStateInfo.blendConstants[3] = 0;

        // 动态参数
        VkPipelineDynamicStateCreateInfo dynamicStateInfo =
        {
            .sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .dynamicStateCount = static_cast<uint32_t>(mPipelineConfig.dynamicState.dynamicStates.size()),
            .pDynamicStates = mPipelineConfig.dynamicState.dynamicStates.data()
        };

        // 图形管线
        VkGraphicsPipelineCreateInfo pipelineInfo =
        {
            .sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .stageCount = ARRAY_SIZE(shaderStageInfo),
            .pStages = shaderStageInfo,
            .pVertexInputState = &vertexInputStateInfo,
            .pInputAssemblyState = &inputAssemblyStateInfo,
            .pTessellationState = nullptr,
            .pViewportState = &viewportStateInfo,
            .pRasterizationState = &rasterizationStateInfo,
            .pMultisampleState = &multisampleStateInfo,
            .pDepthStencilState = &depthStencilStateInfo,
            .pColorBlendState = &colorBlendStateInfo,
            .pDynamicState = &dynamicStateInfo,
            .layout = mPipelineLayout->GetPipelineLayout(),
            .renderPass = mRenderPass->GetRenderPass(),
            .subpass = 0,
            .basePipelineHandle = VK_NULL_HANDLE,
            .basePipelineIndex = 0
        };
        CALL_VK(vkCreateGraphicsPipelines(mDevice->GetLogicDevice(), mDevice->GetPipelineCache(), 1, &pipelineInfo, nullptr, &mPipeline));
    }

    void PiccoloVKPipeline::Bind(VkCommandBuffer commandBuffer)
    {
        vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline);
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetVertexInputState(const std::vector<VkVertexInputBindingDescription> &vertexBindings, const std::vector<VkVertexInputAttributeDescription> &vertexAttributes)
    {
        mPipelineConfig.vertexInputState.vertexBindingDescriptions = vertexBindings;
        mPipelineConfig.vertexInputState.vertexAttributes = vertexAttributes;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetInputAssemblyState(VkPrimitiveTopology topology, VkBool32 primitiveRestartEnable)
    {
        mPipelineConfig.inputAssemblyState.topology = topology;
        mPipelineConfig.inputAssemblyState.primitiveRestartEnable = primitiveRestartEnable;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetRasterizationState(const PipelineRasterizationState &rasterizationState)
    {
        mPipelineConfig.rasterizationState.depthClampEnable = rasterizationState.depthClampEnable;
        mPipelineConfig.rasterizationState.rasterizerDiscardEnable = rasterizationState.rasterizerDiscardEnable;
        mPipelineConfig.rasterizationState.polygonMode = rasterizationState.polygonMode;
        mPipelineConfig.rasterizationState.cullMode = rasterizationState.cullMode;
        mPipelineConfig.rasterizationState.frontFace = rasterizationState.frontFace;
        mPipelineConfig.rasterizationState.depthBiasEnable = rasterizationState.depthBiasEnable;
        mPipelineConfig.rasterizationState.depthBiasConstantFactor = rasterizationState.depthBiasConstantFactor;
        mPipelineConfig.rasterizationState.depthBiasClamp = rasterizationState.depthBiasClamp;
        mPipelineConfig.rasterizationState.depthBiasSlopeFactor = rasterizationState.depthBiasSlopeFactor;
        mPipelineConfig.rasterizationState.lineWidth = rasterizationState.lineWidth;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetMultisampleState(VkSampleCountFlagBits samples, VkBool32 sampleShadingEnable, float minSampleShading)
    {
        mPipelineConfig.multisampleState.rasterizationSamples = samples;
        mPipelineConfig.multisampleState.sampleShadingEnable = sampleShadingEnable;
        mPipelineConfig.multisampleState.minSampleShading = minSampleShading;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetDepthStencilState(const PipelineDepthStencilState &depthStencilState)
    {
        mPipelineConfig.depthStencilState.depthTestEnable = depthStencilState.depthTestEnable;
        mPipelineConfig.depthStencilState.depthWriteEnable = depthStencilState.depthWriteEnable;
        mPipelineConfig.depthStencilState.depthCompareOp = depthStencilState.depthCompareOp;
        mPipelineConfig.depthStencilState.depthBoundsTestEnable = depthStencilState.depthBoundsTestEnable;
        mPipelineConfig.depthStencilState.stencilTestEnable = depthStencilState.stencilTestEnable;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetColorBlendAttachmentState(VkBool32 blendEnable,
        VkBlendFactor srcColorBlendFactor, VkBlendFactor dstColorBlendFactor, VkBlendOp colorBlendOp,
        VkBlendFactor srcAlphaBlendFactor, VkBlendFactor dstAlphaBlendFactor, VkBlendOp alphaBlendOp)
    {
        mPipelineConfig.colorBlendAttachmentState.blendEnable = blendEnable;
        mPipelineConfig.colorBlendAttachmentState.srcColorBlendFactor = srcColorBlendFactor;
        mPipelineConfig.colorBlendAttachmentState.dstColorBlendFactor = dstColorBlendFactor;
        mPipelineConfig.colorBlendAttachmentState.srcAlphaBlendFactor = srcAlphaBlendFactor;
        mPipelineConfig.colorBlendAttachmentState.dstAlphaBlendFactor = dstAlphaBlendFactor;
        mPipelineConfig.colorBlendAttachmentState.alphaBlendOp = alphaBlendOp;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::SetDynamicState(const std::vector<VkDynamicState> &dynamicStates)
    {
        mPipelineConfig.dynamicState.dynamicStates = dynamicStates;
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::EnableAlphaBlend()
    {
        mPipelineConfig.colorBlendAttachmentState =
        {
            .blendEnable = VK_TRUE,
            .srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA,
            .dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
            .colorBlendOp = VK_BLEND_OP_ADD,
            .srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA,
            .dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO,
            .alphaBlendOp = VK_BLEND_OP_ADD
        };
        return this;
    }

    PiccoloVKPipeline* PiccoloVKPipeline::EnableDepthTest()
    {
        mPipelineConfig.depthStencilState.depthTestEnable = VK_TRUE;
        mPipelineConfig.depthStencilState.depthWriteEnable = VK_TRUE;
        mPipelineConfig.depthStencilState.depthCompareOp = VK_COMPARE_OP_LESS;
        return this;
    }
}
